from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter
from langchain.llms import OpenAI
from langchain.document_loaders import(
    DirectoryLoader,
    PagedPDFSplitter,
    PDFMinerLoader
)
from langchain.chains import RetrievalQA

import os
os.environ["OPENAI_API_KEY"] = "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"

loader = DirectoryLoader('./data/', glob="**/*.txt")
documents = loader.load()

text_splitter = CharacterTextSplitter(chunk_size=300, chunk_overlap=0)

split_documents = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()

docsearch = Chroma.from_documents(split_documents, embeddings)


qa = RetrievalQA.from_chain_type(llm=OpenAI(), retriever=docsearch.as_retriever())

results = qa({"query":"什么情况下会拒保？"})

print(results)
